# +----------------------------------------------------------------------
# | ChatWork智能聊天办公系统
# +----------------------------------------------------------------------
# | 软件声明: 本系统并非自由软件,未经授权任何形式的商业使用均属非法。
# | 版权保护: 任何企业和个人不允许对程序代码以任何形式任何目的复制/分发。
# | 授权要求: 如有商业使用需求,请务必先与版权所有者取得联系并获得正式授权。
# +----------------------------------------------------------------------
# | Author: ChatWork Team <2474369941@qq.com>
# +----------------------------------------------------------------------
import os
import json
import time
import logging
from events import scheduler
from common.cache.queue_cache import QueueCache
from common.enums.apply import VideoEnum
from common.enums.wallet import WalletEnum
from common.enums.public import CrontabEnum
from common.utils.cache import RedisUtil
from common.utils.times import TimeUtil
from common.utils.tools import ToolsUtil
from common.utils.urls import UrlUtil
from common.models.sys import SysCrontabModel
from common.models.users import UserWalletModel
from common.models.video import VideoRecordsModel
from common.chain.video_server import VideoZhipuServer
from common.chain.video_server import VideoOpenAiHkServer
from plugins.storage.driver import StorageDriver


logger = logging.getLogger(__name__)


async def execute(**kwargs):
    """ AI视频回调任务 """
    start_time = time.time()
    crontab_id: int = int(kwargs["w_id"])
    crontab_job: str = str(kwargs["w_job"])
    if not scheduler.get_job(crontab_job):
        scheduler.pause_job(crontab_job)

    try:
        record_id = await QueueCache.pop_queue_value(QueueCache.VIDEO_JOB)
        if not record_id or record_id is None:
            return True

        await dispense(int(record_id))

        await SysCrontabModel.compute(crontab_id, start_time)
    except Exception as e:
        await SysCrontabModel.compute(crontab_id, start_time, status=CrontabEnum.CRON_ERROR, error=str(e))


async def dispense(record_id: int):
    # 查询任务
    record = await VideoRecordsModel.filter(id=record_id).first()
    if not record and record.status != VideoEnum.STATUS_ING:
        return

    # 处理响应
    if record.engine == "luma":
        await handle_luma_check(record)
    elif record.engine == "zhipu":
        await handle_zhipu_check(record)
    else:
        record.status = VideoEnum.STATUS_FAIL
        record.fail_reason = f"No processing channel {record.engine}"
        await record.save()


async def handle_luma_check(record):
    """
    Luma视频回调处理

    Document:
        https://www.volcengine.com/docs/6791/1366783

    Author:
        zero
    """
    cover_dir_path = ""
    video_dir_path = ""
    response = {}
    try:
        # 限流拦截
        lock_key = f"video:prevent_{record.id}"
        lock_cache = await RedisUtil.get(lock_key)
        if lock_cache:
            print("已锁定: " + str(record.id))
            await QueueCache.push_video(record.id)
            return False

        # 锁定记录
        await RedisUtil.set(lock_key, record.id, 60)
        print(f"Luma视频开始: [重试={record.retry_count}, 任务={record.id}, 时间={time.time()}")

        # 获取生成结果
        server = VideoOpenAiHkServer()
        response = await server.fetch(record.task_id)

        # 获取结果异常
        if response.get("error") or not response.get("data"):
            error = response.get("error") or {}
            await VideoRecordsModel.filter(id=record.id).update(
                status=int(VideoEnum.STATUS_FAIL),
                fail_reason=error.get("message", "未知错误"),
                notify_snap=json.dumps(response, ensure_ascii=False),
                update_time=int(time.time())
            )
            await UserWalletModel.inc(
                scene=WalletEnum.POINTS,
                user_id=record.user_id,
                change_type=WalletEnum.POINTS_INC_VIDEO_FAIL,
                change_amount=record.use_points,
                project="AI视频",
                remarks=error.get("message", "未知错误")
            )
            return None

        # 处理生成结果
        data = response.get("data")
        if data[0].get("state") == "completed":
            # 成功
            result = data[0]["artifact"]
            cover_url: str = result["last_frame"]["url"]
            video_url: str = result["video"]["url"]

            # 文件扩展
            imageExt: str = str(cover_url.split(".")[-1]).strip().lower()
            videoExt: str = str(video_url.split(".")[-1]).strip().lower()

            # 存储路径
            filename: str = ToolsUtil.make_md5_str(str(video_url)+str(time.time()))
            curr_now: str = TimeUtil.timestamp_to_date(int(time.time()), str("%Y%m%d"))
            cover_dir: str = f"storage/applies/video/{curr_now}/{filename}.{imageExt}"
            video_dir: str = f"storage/applies/video/{curr_now}/{filename}.{videoExt}"
            cover_dir_path: str = UrlUtil.to_root_path(path=cover_dir)
            video_dir_path: str = UrlUtil.to_root_path(path=video_dir)

            # 下载封面
            await ToolsUtil.download_file(url=cover_url, save_path=cover_dir_path)

            # 下载视频
            await ToolsUtil.download_file(url=video_url, save_path=video_dir_path)

            # 视频上传OSS
            engine: str = await UrlUtil.get_storage_engine()
            if engine != "local":
                await StorageDriver.push(path=cover_dir_path, key=cover_dir)
                await StorageDriver.push(path=video_dir_path, key=video_dir)

            # 更新记录
            await VideoRecordsModel.filter(id=record.id).update(
                status=VideoEnum.STATUS_YES,
                cover_url=str(cover_dir),
                video_url=str(video_dir),
                fail_reason="",
                notify_snap=""
            )

            # 生成成功
            task_time = int(time.time()) - int(record.start_time)
            print(f"Luma视频成功: [重试={record.retry_count}, 任务={record.id}, 耗时={task_time}s")
        elif data[0].get("state") == "PROCESSING":
            # 进行中
            await QueueCache.push_video(record.id)

    except Exception as e:
        if record.retry_count < 2:
            print(f"Luma视频重试: [重试={record.retry_count + 1}, 任务={record.id}, 错误={str(e)}")
            record.retry_count += 1
            record.start_time = int(time.time())
            await record.save()
            await QueueCache.push_video(record.id)
        else:
            logger.error("Luma视频回调失败: [%s]", str(e))
            await VideoRecordsModel.filter(id=record.id).update(
                status=VideoEnum.STATUS_FAIL,
                notify_snap=json.dumps(response, ensure_ascii=False),
                fail_reason=str(e),
                update_time=int(time.time())
            )
            if record.use_points > 0:
                await UserWalletModel.inc(
                    scene=WalletEnum.POINTS,
                    user_id=record.user_id,
                    change_type=WalletEnum.POINTS_INC_VIDEO_FAIL,
                    change_amount=record.use_points,
                    project="AI视频",
                    remarks=str(e)
                )
    finally:
        engine: str = await UrlUtil.get_storage_engine()
        if engine != "local":
            if cover_dir_path and os.path.exists(cover_dir_path):
                os.unlink(cover_dir_path)
            if video_dir_path and os.path.exists(video_dir_path):
                os.unlink(video_dir_path)


async def handle_zhipu_check(record):
    """
    智谱视频回调处理

    Document:
        https://www.volcengine.com/docs/6791/1366783
    """
    cover_dir_path = ""
    video_dir_path = ""
    response = {}
    try:
        # 限流拦截
        lock_key = f"video:prevent_{record.id}"
        lock_cache = await RedisUtil.get(lock_key)
        if lock_cache:
            await QueueCache.push_video(record.id)
            return False

        # 锁定记录
        await RedisUtil.set(lock_key, record.id, 70)
        print(f"智谱视频开始: [重试={record.retry_count}, 任务={record.id}, 时间={time.time()}]")

        # 获取生成结果
        server = VideoZhipuServer()
        response = await server.fetch(record.task_id)

        # 根据状态处理: [PROCESSING=处理中, SUCCESS=成功, FAIL=失败]
        if response.get("task_status") == "SUCCESS":
            result = response.get("video_result")[0]
            cover_url: str = result.get("cover_image_url")
            video_url: str = result.get("url")

            # 文件扩展
            imageExt: str = str(cover_url.split(".")[-1]).strip().lower()
            videoExt: str = str(video_url.split(".")[-1]).strip().lower()

            # 存储路径
            filename: str = ToolsUtil.make_md5_str(video_url)
            curr_now: str = TimeUtil.timestamp_to_date(int(time.time()), str("%Y%m%d"))
            cover_dir: str = f"storage/applies/video/{curr_now}/{filename}.{imageExt}"
            video_dir: str = f"storage/applies/video/{curr_now}/{filename}.{videoExt}"
            cover_dir_path: str = UrlUtil.to_root_path(path=str(cover_dir))
            video_dir_path: str = UrlUtil.to_root_path(path=str(video_dir))

            # 下载封面
            await ToolsUtil.download_file(url=cover_url, save_path=cover_dir_path)

            # 下载视频
            await ToolsUtil.download_file(url=video_url, save_path=video_dir_path)

            # 视频上传OSS
            engine: str = await UrlUtil.get_storage_engine()
            if engine != str("local"):
                await StorageDriver.push(path=cover_dir_path, key=cover_dir)
                await StorageDriver.push(path=video_dir_path, key=video_dir)

            # 更新记录
            await VideoRecordsModel.filter(id=record.id).update(
                cover_url=str(cover_dir),
                video_url=str(video_dir),
                status=VideoEnum.STATUS_YES
            )

            # 生成成功
            task_time = int(time.time()) - int(record.start_time)
            print(f"智谱视频成功: [重试={record.retry_count}, 任务={record.id}, 耗时={task_time}s]")
        elif response.get("task_status") == "PROCESSING":
            await QueueCache.push_video(record.id)
        else:
            error = response.get("error") or {}
            message = str(error.get("code", "0")) + ":" + error.get("message", "未知错误")
            await VideoRecordsModel.filter(id=record.id).update(
                status=VideoEnum.STATUS_FAIL,
                fail_reason=message,
                notify_snap=json.dumps(response, ensure_ascii=False),
                update_time=int(time.time())
            )
            await UserWalletModel.inc(
                scene=WalletEnum.POINTS,
                user_id=record.user_id,
                change_type=WalletEnum.POINTS_INC_VIDEO_FAIL,
                change_amount=record.use_points,
                project="AI视频",
                remarks=message
            )
    except Exception as e:
        if record.retry_count < 2:
            print(f"智谱视频重试: [重试={record.retry_count + 1}, 任务={record.id}, 错误={str(e)}]")
            record.retry_count += 1
            record.start_time = int(time.time())
            await record.save()
            await QueueCache.push_video(record.id)
        else:
            logger.error("AI智谱视频回调失败: [%s]", str(e))
            await VideoRecordsModel.filter(id=record.id).update(
                status=VideoEnum.STATUS_FAIL,
                fail_reason=str(e),
                notify_snap=json.dumps(response, ensure_ascii=False),
                update_time=int(time.time())
            )
            if record.use_points > 0:
                await UserWalletModel.inc(
                    scene=WalletEnum.POINTS,
                    user_id=record.user_id,
                    change_type=WalletEnum.POINTS_INC_VIDEO_FAIL,
                    change_amount=record.use_points,
                    project="AI视频",
                    remarks=str(e)
                )
    finally:
        engine: str = await UrlUtil.get_storage_engine()
        if engine != "local":
            if cover_dir_path and os.path.exists(cover_dir_path):
                os.unlink(cover_dir_path)
            if video_dir_path and os.path.exists(video_dir_path):
                os.unlink(video_dir_path)
